# RTSP流处理模块
import cv2
import threading
import time
import logging
from typing import Optional, Callable, Dict, Any
from queue import Queue
import numpy as np


class RTSPStreamProcessor:
    """RTSP流处理器"""

    def __init__(self, rtsp_url: str, app_config: Dict[str, Any]):
        """
        初始化流处理器
        Args:
            rtsp_url: RTSP流地址
            app_config: 应用配置
        """
        self.rtsp_url = rtsp_url
        self.config = app_config
        self.cap = None
        self.running = False
        self.frame_queue = Queue(maxsize=5)  # 限制队列大小以减少延迟
        self.current_frame = None
        self.frame_lock = threading.Lock()

        # 跳帧计数器
        self.frame_counter = 0
        self.skip_frames = app_config['detection']['skip_frames']

        logging.info(f"RTSPStreamProcessor初始化，URL: {rtsp_url}")

    def connect(self) -> bool:
        """
        连接到RTSP流
        Returns:
            连接是否成功
        """
        try:
            self.cap = cv2.VideoCapture(self.rtsp_url)

            # 设置缓冲区大小以减少延迟
            self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)

            if not self.cap.isOpened():
                logging.error("无法打开RTSP流")
                return False

            logging.info("成功连接到RTSP流")
            return True

        except Exception as e:
            logging.error(f"连接RTSP流时出错: {e}")
            return False

    def start_capture(self, frame_callback: Optional[Callable] = None):
        """
        开始捕获视频流
        Args:
            frame_callback: 帧回调函数
        """
        if not self.connect():
            return

        self.running = True
        capture_thread = threading.Thread(
            target=self._capture_loop,
            args=(frame_callback,)
        )
        capture_thread.daemon = True
        capture_thread.start()

        logging.info("开始捕获视频流")

    def _capture_loop(self, frame_callback: Optional[Callable] = None):
        """
        捕获循环
        Args:
            frame_callback: 帧回调函数
        """
        while self.running and self.cap is not None:
            try:
                ret, frame = self.cap.read()

                if not ret:
                    logging.warning("无法读取帧，尝试重连...")
                    time.sleep(1)
                    if not self._reconnect():
                        break
                    continue

                # 跳帧处理以减少延迟
                self.frame_counter += 1
                if self.frame_counter % (self.skip_frames + 1) != 0:
                    continue

                # 更新当前帧
                with self.frame_lock:
                    self.current_frame = frame.copy()

                # 调用回调函数
                if frame_callback:
                    frame_callback(frame)

                # 将帧放入队列
                if not self.frame_queue.full():
                    self.frame_queue.put(frame)

            except Exception as e:
                logging.error(f"捕获帧时出错: {e}")
                time.sleep(0.1)

    def _reconnect(self) -> bool:
        """
        重新连接RTSP流
        Returns:
            重连是否成功
        """
        try:
            if self.cap:
                self.cap.release()

            # 等待一段时间再重连
            time.sleep(self.config['rtsp']['reconnect_delay'])

            return self.connect()

        except Exception as e:
            logging.error(f"重连时出错: {e}")
            return False

    def get_latest_frame(self) -> Optional[np.ndarray]:
        """
        获取最新帧
        Returns:
            最新的帧，如果没有则返回None
        """
        with self.frame_lock:
            return self.current_frame.copy() if self.current_frame is not None else None

    def stop(self):
        """停止捕获"""
        self.running = False
        if self.cap:
            self.cap.release()

        # 清空队列
        while not self.frame_queue.empty():
            try:
                self.frame_queue.get_nowait()
            except:
                break

        logging.info("视频流捕获已停止")
